#include "h264_encode.hpp"
#include <glog/logging.h>
#include <exception>
H264Encoder::H264Encoder(int image_widht, int image_height, unsigned int bitrate, int fps)
    : codec_ctx_(nullptr), codec_(nullptr), frame_(nullptr), pkt_(nullptr), image_width_(image_widht), image_height_(image_height), frame_bytes_(0), yuv_buffer_(nullptr), index_(0)
{
    codec_ = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (NULL == codec_)
    {
        LOG(INFO) << "find AV_CODEC_ID_H264 encoder fail";
        throw std::bad_exception();
    }
    codec_ctx_ = avcodec_alloc_context3(codec_);
    if (NULL == codec_ctx_)
    {
        LOG(INFO) << "avcodec_alloc_context3 failed";
        throw std::bad_exception();
    }
    codec_ctx_->width = image_width_;
    codec_ctx_->height = image_height_;
    codec_ctx_->bit_rate = bitrate;
    AVRational r = {1, fps};
    codec_ctx_->time_base = r;
    codec_ctx_->framerate = (AVRational){fps, 1};
    codec_ctx_->gop_size = 12;
    codec_ctx_->pix_fmt = AV_PIX_FMT_YUV420P;
    // av_opt_set(codec_ctx_->priv_data,"preset","slow",0);
    av_opt_set(codec_ctx_->priv_data, "preset", "medium", 0);
    if (avcodec_open2(codec_ctx_, codec_, NULL) < 0)
    {
        LOG(INFO) << "avcodec_open2 error";
        throw std::bad_exception();
    }
    frame_ = av_frame_alloc();
    if (NULL == frame_)
    {
        LOG(INFO) << "av_frame_alloc failed";
        throw std::bad_exception();
    }
    //分配AVFrame及像素存储空间
    frame_->width = image_width_;
    frame_->height = image_height_;
    frame_->format = codec_ctx_->pix_fmt;
    if (av_frame_get_buffer(frame_, 0) < 0)
    {
        LOG(INFO) << "av_frame_get_buffer failed";
        throw std::bad_exception();
    }
    //分配pkt
    pkt_ = av_packet_alloc();
    if (NULL == pkt_)
    {
        LOG(INFO) << "av_packet_alloc failed";
        throw std::bad_exception();
    }
    frame_bytes_ = av_image_get_buffer_size(codec_ctx_->pix_fmt, frame_->width, frame_->height, 1);
    yuv_buffer_ = new uint8_t[frame_bytes_];
}

H264Encoder::~H264Encoder()
{
    av_frame_free(&frame_);
    av_packet_free(&pkt_);
    avcodec_free_context(&codec_ctx_);
    delete[] yuv_buffer_;
}
/*
int get_nalu_index(uint8_t *buffer, size_t buffer_len, size_t offset)
{
    size_t i = offset + 2;
    while((++i)<buffer_len)
    {
        if(buffer[i-3]== 0x00 && buffer[i-2]==0x00 && buffer[i-1]==0x01)
        {
            int
        }
    }
}
*/
int H264Encoder::encode(uint8_t *image, uint8_t *h264_buffer, size_t *buffer_size)
{
    if (image == nullptr || h264_buffer == nullptr || buffer_size == nullptr || *buffer_size == 0)
    {
        return -1;
    }
    index_+=40;
    av_image_fill_arrays(frame_->data, frame_->linesize, image, (AVPixelFormat)frame_->format, frame_->width, frame_->height, 1);
    frame_->pts = index_;
    
    int ret = avcodec_send_frame(codec_ctx_, frame_);

    if (ret < 0)
    {
        LOG(INFO) << "avcodec_send_frame failed";
        return -1;
    }
    ret = avcodec_receive_packet(codec_ctx_, pkt_);
    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
    {
        return 0;
    }
    if (ret < 0)
    {
        LOG(INFO) << "avcodec_receive_packet failed";
        return -1;
    }
    size_t total_size = pkt_->size;
    if (total_size > *buffer_size)
    {
        return -1;
    }
    else
    {
        memcpy(h264_buffer, pkt_->data, total_size);
        *buffer_size = total_size;
        frame_->pts++;
        return 0;
    }
}
